// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/constant_expr.h>

namespace kumo::pollux::exec {
    void ConstantExpr::evalSpecialForm(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        if (sharedConstantValue_.use_count() == 1) {
            sharedConstantValue_->resize(rows.end());
        } else {
            // By reassigning sharedConstantValue_ we increase the chances that it will
            // be unique the next time this expression is evaluated.
            sharedConstantValue_ =
                    BaseVector::wrap_in_constant(rows.end(), 0, sharedConstantValue_);
        }

        if (needToSetIsAscii_) {
            // sharedConstantValue_ must be unique because computeAndSetIsAscii may
            // modify it.
            POLLUX_CHECK_EQ(sharedConstantValue_.use_count(), 1);
            auto *vector =
                    sharedConstantValue_->as_unchecked<SimpleVector<StringView> >();
            LocalSingleRow singleRow(context, 0);
            bool isAscii = vector->computeAndSetIsAscii(*singleRow);
            vector->setAllIsAscii(isAscii);
            needToSetIsAscii_ = false;
        }

        context.moveOrCopyResult(sharedConstantValue_, rows, result);
    }

    void ConstantExpr::evalSpecialFormSimplified(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        ExceptionContextSetter exceptionContext(
            {
                [](PolluxException::Type /*exceptionType*/, auto *expr) {
                    return static_cast<Expr *>(expr)->toString();
                },
                this
            });

        // Simplified path should never ask us to write to a vector that was already
        // pre-allocated.
        POLLUX_CHECK_NULL(result);

        result = BaseVector::wrap_in_constant(rows.end(), 0, sharedConstantValue_);
    }

    std::string ConstantExpr::toString(bool /*recursive*/) const {
        return fmt::format(
            "{}:{}", sharedConstantValue_->toString(0), type()->toString());
    }

    namespace {
        void appendSqlLiteral(
            const BaseVector &vector,
            vector_size_t row,
            std::ostream &out);

        void appendSqlLiteralList(
            const BaseVector &vector,
            vector_size_t offset,
            vector_size_t size,
            std::ostream &out) {
            for (auto i = offset; i < offset + size; ++i) {
                if (i > offset) {
                    out << ", ";
                }
                appendSqlLiteral(vector, i, out);
            }
        }

        void appendSqlString(const std::string &value, std::ostream &out) {
            // Escape single quotes: ' -> ''.
            static constexpr char kSingleQuote = '\'';

            out << kSingleQuote;
            auto prevPos = 0;
            auto pos = value.find(kSingleQuote);
            while (pos != std::string::npos) {
                out << value.substr(prevPos, pos - prevPos) << kSingleQuote;
                prevPos = pos;
                pos = value.find(kSingleQuote, prevPos + 1);
            }
            out << value.substr(prevPos, pos) << kSingleQuote;
            return;
        }

        std::string toSqlType(const TypePtr &type) {
            switch (type->kind()) {
                case TypeKind::ARRAY:
                    return fmt::format("{}[]", toSqlType(type->childAt(0)));
                case TypeKind::MAP:
                    return fmt::format(
                        "MAP({}, {})",
                        toSqlType(type->childAt(0)),
                        toSqlType(type->childAt(1)));
                case TypeKind::ROW: {
                    std::ostringstream out;
                    out << "STRUCT(";
                    for (auto i = 0; i < type->size(); ++i) {
                        if (i > 0) {
                            out << ", ";
                        }
                        out << type->as_row().nameOf(i) << " " << toSqlType(type->childAt(i));
                    }
                    out << ")";
                    return out.str();
                }
                default:
                    return type->toString();
            }
        }

        void appendSqlLiteral(
            const BaseVector &vector,
            vector_size_t row,
            std::ostream &out) {
            if (vector.is_null_at(row)) {
                if (vector.type()->containsUnknown()) {
                    out << "NULL";
                } else {
                    out << "NULL::" << toSqlType(vector.type());
                }
                return;
            }

            switch (vector.type_kind()) {
                case TypeKind::BOOLEAN: {
                    auto value = vector.as<SimpleVector<bool> >()->value_at(row);
                    out << (value ? "TRUE" : "FALSE");
                    break;
                }
                case TypeKind::INTEGER: {
                    if (vector.type()->isDate()) {
                        auto dateVector = vector.wrapped_vector()->as<SimpleVector<int32_t> >();
                        out << "'"
                                << DATE()->toString(dateVector->value_at(vector.wrapped_index(row)))
                                << "'::" << vector.type()->toString();
                    } else {
                        out << "'" << vector.wrapped_vector()->toString(vector.wrapped_index(row))
                                << "'::" << vector.type()->toString();
                    }
                    break;
                }
                case TypeKind::TINYINT:
                    [[fallthrough]];
                case TypeKind::SMALLINT:
                    [[fallthrough]];
                case TypeKind::BIGINT: {
                    if (vector.type()->isIntervalDayTime()) {
                        auto *intervalVector =
                                vector.wrapped_vector()->as<SimpleVector<int64_t> >();
                        out << "INTERVAL " << intervalVector->value_at(vector.wrapped_index(row))
                                << " MILLISECOND";
                        break;
                    }
                    [[fallthrough]];
                }
                case TypeKind::HUGEINT:
                    [[fallthrough]];
                case TypeKind::REAL:
                    [[fallthrough]];
                case TypeKind::DOUBLE:
                    out << "'" << vector.wrapped_vector()->toString(vector.wrapped_index(row))
                            << "'::" << vector.type()->toString();
                    break;
                case TypeKind::TIMESTAMP: {
                    TimestampToStringOptions options;
                    options.dateTimeSeparator = ' ';
                    const auto ts =
                            vector.wrapped_vector()->as<SimpleVector<Timestamp> >()->value_at(row);
                    out << "'" << ts.toString(options) << "'::" << vector.type()->toString();
                    break;
                }
                case TypeKind::VARCHAR:
                    appendSqlString(
                        vector.wrapped_vector()->toString(vector.wrapped_index(row)), out);
                    break;
                case TypeKind::ARRAY: {
                    out << "ARRAY[";
                    auto array_vector = vector.wrapped_vector()->as<ArrayVector>();
                    auto arrayRow = vector.wrapped_index(row);
                    auto offset = array_vector->offsetAt(arrayRow);
                    auto size = array_vector->sizeAt(arrayRow);
                    appendSqlLiteralList(*array_vector->elements(), offset, size, out);
                    out << "]";
                    break;
                }
                case TypeKind::MAP: {
                    out << "map(ARRAY[";
                    auto map_vector = vector.wrapped_vector()->as<MapVector>();
                    auto mapRow = vector.wrapped_index(row);
                    auto offset = map_vector->offsetAt(mapRow);
                    auto size = map_vector->sizeAt(mapRow);
                    appendSqlLiteralList(*map_vector->mapKeys(), offset, size, out);
                    out << "], ARRAY[";
                    appendSqlLiteralList(*map_vector->mapValues(), offset, size, out);
                    out << "])";
                    break;
                }
                case TypeKind::ROW: {
                    out << "row_constructor(";
                    auto row_vector = vector.wrapped_vector()->as<RowVector>();
                    auto baseRow = vector.wrapped_index(row);
                    for (auto i = 0; i < row_vector->childrenSize(); ++i) {
                        if (i > 0) {
                            out << ", ";
                        }
                        appendSqlLiteral(*row_vector->childAt(i), baseRow, out);
                    }
                    out << ")";
                    break;
                }
                default:
                    // TODO: update ExprStatsTest.exceptionPreparingStatsForListener once
                    // support for VARBINARY is added.
                    POLLUX_UNSUPPORTED(
                        "Type not supported yet: {}", vector.type()->toString());
            }
        }

        bool canBeExpressedInSQL(const TypePtr &type) {
            // Logical types cannot be expressed in SQL.
            const bool isLogicalType = type->name() != type->kindName();
            return type->isPrimitiveType() && type != VARBINARY() && !isLogicalType;
        }
    } // namespace

    std::string ConstantExpr::toSql(
        std::vector<VectorPtr> *complexConstants) const {
        POLLUX_CHECK_NOT_NULL(sharedConstantValue_);

        // TODO canBeExpressedInSQL is misleading. ARRAY(INTEGER()) can be expressed
        // in SQL, but canBeExpressedInSQL returns false for that type. we need to
        // distinguish between types that cannot be expressed in SQL at all
        // (VARBINARY, logic types like JSON) and types that can be expressed in
        // SQL, but they can be large and therefore we want to provider an option to
        // represent then using 'complex constants'. If a type cannot be expressed
        // in SQL, we should assert that complexConstants is not null.

        std::ostringstream out;
        if (complexConstants && !canBeExpressedInSQL(sharedConstantValue_->type())) {
            int idx = complexConstants->size();
            out << "__complex_constant(c" << idx << ")";
            complexConstants->push_back(sharedConstantValue_);
        } else {
            appendSqlLiteral(*sharedConstantValue_, 0, out);
        }
        return out.str();
    }
} // namespace kumo::pollux::exec
